# 这个脚本主要用来实现基本的模型

import os
from dotenv import load_dotenv 
from openai import OpenAI

# 模型
from langchain.chat_models import init_chat_model

# 链式调用方法
from langchain_core.output_parsers import StrOutputParser
from langchain.chat_models import init_chat_model

# 添加提示词模板
from langchain.output_parsers.boolean import BooleanOutputParser
from langchain.prompts import ChatPromptTemplate

# 结构化输出的模板
from langchain_core.prompts import PromptTemplate
from langchain.output_parsers import ResponseSchema, StructuredOutputParser

# 复合链的库
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnableSequence
from langchain.output_parsers import ResponseSchema, StructuredOutputParser

# 自定义可运行环节的库
from langchain_core.runnables import RunnableLambda

# 构建流式智能问答系统
from langchain_core.output_parsers import StrOutputParser
from langchain.chat_models import init_chat_model
from langchain.prompts import ChatPromptTemplate


class basic_mode_project:
    def __init__(self) -> None:
        load_dotenv(override=True)
        DeepSeek_API_KEY = os.getenv("DEEPSEEK_API_KEY")
        self.ds_api=DeepSeek_API_KEY

        #print(DeepSeek_API_KEY)  # 可以通过打印查看
        self.model = init_chat_model(model="deepseek-chat", model_provider="deepseek")  
        
    
    def most_basic_model(self):
        client = OpenAI(api_key=self.ds_api, base_url="https://api.deepseek.com")
        # 调用DeepSeek的API，生成回答
        response = client.chat.completions.create(
            model="deepseek-chat",
            messages=[
                {"role": "system", "content": "你是乐于助人的助手，请根据用户的问题给出回答"},
                {"role": "user", "content": "你好，请你介绍一下你自己。"},
            ],
        )
        # 打印模型最终的响应结果
        print(response.choices[0].message.content)

    def basic_chatbot_model(self):
        # model可以选择deepseek-chat, deepseek-reasoner
        model = init_chat_model(model="deepseek-chat", model_provider="deepseek")  
        question = "你好，请你介绍一下你自己。"
        result = model.invoke(question)
        print(result.content)
        print(result.additional_kwargs)
        '''
        deepseek-chat的result的内容
        AIMessage(
        content='你好！我是 **DeepSeek Chat**，由深度求索公司（DeepSeek）研发的一款智能AI助手。我可以帮助你解答各种问题，包括学习、工作、编程、写作、生活百科等多个领域。  \n\n### **我的特点：**  \n✅ **免费使用**：目前无需付费，你可以随时向我提问！  \n✅ **知识丰富**：我的知识截止到 **2024年7月**，可以为你提供较新的信息。  \n✅ **超长上下文支持**：可以处理 **128K** 长度的文本，适合分析长文档或复杂问题。  \n✅ **文件阅读**：支持上传 **PDF、Word、Excel、PPT、TXT** 等文件，并从中提取信息进行分析。  \n✅ **多语言能力**：可以用中文、英文等多种语言交流，帮助你翻译或学习外语。  \n✅ **编程助手**：能写代码、调试、优化算法，支持Python、C++、Java等多种编程语言。  \n\n### **我能帮你做什么？**  \n📖 **学习辅导**：解题思路、论文写作、知识点讲解  \n💼 **工作效率**：写邮件、做PPT、总结报告  \n💡 **创意灵感**：写故事、起名字、头脑风暴  \n📊 **数据分析**：处理表格、绘制图表、解读数据  \n🔧 **技术支持**：代码调试、算法优化、技术咨询  \n\n你可以随时向我提问，我会尽力提供最准确、有用的回答！😊 有什么我可以帮你的吗？', 
        additional_kwargs={'refusal': None}, 
        response_metadata={'token_usage': {'completion_tokens': 311, 'prompt_tokens': 9, 'total_tokens': 320, 'completion_tokens_details': None, 'prompt_tokens_details': {'audio_tokens': None, 'cached_tokens': 0}, 'prompt_cache_hit_tokens': 0, 'prompt_cache_miss_tokens': 9}, 'model_name': 'deepseek-chat', 'system_fingerprint': 'fp_8802369eaa_prod0425fp8', 'id': 'e3323dc0-46fa-497b-aaf5-83596184f8b2', 'service_tier': None, 'finish_reason': 'stop', 'logprobs': None}, 
        id='run--9f652093-6feb-484f-99f7-0c30e33290ab-0', 
        usage_metadata={'input_tokens': 9, 'output_tokens': 311, 'total_tokens': 320, 'input_token_details': {'cache_read': 0}, 'output_token_details': {}}
        )
        '''


    def basic_chain_model(self):
        # 这是一个基础的链式调用

        # 使用deepseek模型
        model = init_chat_model(model="deepseek-chat", model_provider="deepseek")  

        #一个简单的链式调用
        # 直接使用模型 + 输出解析器搭建一个链
        basic_qa_chain = model | StrOutputParser()  #这里的输出解析器是指将输出的结果仅保留为字符串
        question = "你好，请你介绍一下你自己。"
        # 链式调用
        result=basic_qa_chain.invoke(question)
        print(result)

    def prompt_template_model(self):
        # 这是一个单个的链式调用的应用
        # 这是提示词的模板

        # 使用deepseek模型
        model = init_chat_model(model="deepseek-chat", model_provider="deepseek")  

        prompt_template = ChatPromptTemplate([
            ("system", "你是一个乐意助人的助手，请根据用户的问题给出回答"),
            ("user", "这是用户的问题： {topic}， 请用 yes 或 no 来回答")
        ])

        # 直接使用模型 + 输出解析器
        bool_qa_chain = prompt_template | model | StrOutputParser() # 输出yes or no
        # bool_qa_chain = prompt_template | model | BooleanOutputParser() # 输出 True or False

        # 测试
        question = "请问 1 + 1 是否 大于 2？"
        result = bool_qa_chain.invoke(question)
        print(result)

    def tructured_output_model(self):
        # 这是一个单个的链式调用的应用
        # 这是结构化输出

        model = init_chat_model(model="deepseek-chat", model_provider="deepseek")          

        schemas = [
            ResponseSchema(name="name", description="用户的姓名"),
            ResponseSchema(name="age", description="用户的年龄")
        ]
        parser = StructuredOutputParser.from_response_schemas(schemas)

        prompt = PromptTemplate.from_template(
            "请根据以下内容提取用户信息，并返回 JSON 格式：\n{input}\n\n{format_instructions}"
        )

        chain = (
            prompt.partial(format_instructions=parser.get_format_instructions())
            | model
            | parser
        )

        result = chain.invoke({"input": "用户叫李雷，今年25岁，是一名工程师。"})
        print(result)  

    def multi_chain_model(self):

        model = init_chat_model(model="deepseek-chat", model_provider="deepseek")   
        # 第一步：根据标题生成新闻正文
        news_gen_prompt = PromptTemplate.from_template(
            "请根据以下新闻标题撰写一段简短的新闻内容（100字以内）：\n\n标题：{title}"
        )

        # 第一个子链：生成新闻内容
        news_chain = news_gen_prompt | model

        # 第二步：从正文中提取结构化字段
        schemas = [
            ResponseSchema(name="time", description="事件发生的时间"),
            ResponseSchema(name="location", description="事件发生的地点"),
            ResponseSchema(name="event", description="发生的具体事件"),
        ]
        parser = StructuredOutputParser.from_response_schemas(schemas)

        summary_prompt = PromptTemplate.from_template(
            "请从下面这段新闻内容中提取关键信息，并返回结构化JSON格式：\n\n{news}\n\n{format_instructions}"
        )

        # 第二个子链：生成新闻摘要
        summary_chain = (
            summary_prompt.partial(format_instructions=parser.get_format_instructions())
            | model
            | parser
        )

        # 组合成一个复合 Chain
        full_chain = news_chain | summary_chain

        # 调用复合链
        result = full_chain.invoke({"title": "苹果公司在加州发布新款AI芯片"})
        print(result)

    def self_define_multi_chain_model(self, model):
        # 第一步：根据标题生成新闻正文
        news_gen_prompt = PromptTemplate.from_template(
            "请根据以下新闻标题撰写一段简短的新闻内容（100字以内）：\n\n标题：{title}"
        )

        # 第一个子链：生成新闻内容
        news_chain = news_gen_prompt | model


        # 一个简单的打印函数，调试用
        def debug_print(x):
            print("中间结果（新闻正文）:", x)
            return x

        debug_node = RunnableLambda(debug_print) # 设置自定义的可运行的节点

        # 第二步：从正文中提取结构化字段
        schemas = [
            ResponseSchema(name="time", description="事件发生的时间"),
            ResponseSchema(name="location", description="事件发生的地点"),
            ResponseSchema(name="event", description="发生的具体事件"),
        ]
        parser = StructuredOutputParser.from_response_schemas(schemas)

        summary_prompt = PromptTemplate.from_template(
            "请从下面这段新闻内容中提取关键信息，并返回结构化JSON格式：\n\n{news}\n\n{format_instructions}"
        )

        # 第二个子链：生成新闻摘要
        summary_chain = (
            summary_prompt.partial(format_instructions=parser.get_format_instructions())
            | model
            | parser
        )

        # 插入 debug 节点
        full_chain = news_chain | debug_node | summary_chain

        # 调用复合链
        result = full_chain.invoke({"title": "苹果公司在加州发布新款AI芯片"})
        print(result)

    async def ai_service_model(self, model):
        # TODO: 异步模型的调用目前不行
        chatbot_prompt = ChatPromptTemplate.from_messages([
            ("system", "你叫小智，是一名乐于助人的助手。"),
            ("user", "{input}")
        ])

        # 使用 DeepSeek 模型
        model = init_chat_model(model="deepseek-chat", model_provider="deepseek")  

        # 直接使用模型 + 输出解析器
        #basic_qa_chain = chatbot_prompt | model | StrOutputParser()
        qa_chain_with_system = chatbot_prompt | model | StrOutputParser()

        # 测试：简单的问答
        '''
        question = "你好，请你介绍一下你自己。"
        result = basic_qa_chain.invoke(question)
        print(result)
        '''

        # 异步实现流式输出
        async for chunk in qa_chain_with_system.astream({"input": "你好，请你介绍一下你自己"}):
            print(chunk, end="", flush=True)

if __name__=='__main__':
    model_test=basic_mode_project()
    model_test.ai_service_model(model_test.model)







